import os
import pandas as pd
import re

# 定义文件夹路径
RAW_DATA_FOLDER = 'raw_data/'
CLEAN_DATA_FOLDER = 'clean_data/'

# 创建清洗后数据的文件夹（如果不存在）
os.makedirs(CLEAN_DATA_FOLDER, exist_ok=True)

# 根据文本内容校准标签的函数
def correct_labels(row_text):
    # 定义塔台指令的关键词和语言特征
    tower_keywords = (
        r"塔台|请|指令|跑道|起飞|放行|风速|气压|天气|"
        r"爬升至|离场|航路|允许|批准|转向|报告|校准|下降|滑行|等待"
    )
    
    # 定义飞行员响应的关键词和语言特征
    pilot_keywords = (
        r"收到|明白|执行|确认|复述|报告|已转向|正在|准备|完成|达到|遵循|复制|收到指令"
    )
    
    # 判断标签
    if re.search(tower_keywords, row_text):
        return "塔台"
    elif re.search(pilot_keywords, row_text):
        return "飞行员"
    else:
        return "飞行员"  # 默认视为飞行员

# 处理文件夹中的所有Excel文件
for filename in os.listdir(RAW_DATA_FOLDER):
    if filename.endswith('.xlsx'):
        # 读取Excel文件
        file_path = os.path.join(RAW_DATA_FOLDER, filename)
        df = pd.read_excel(file_path, header=None, names=["时间", "文本内容", "场景", "标签"])
        
        # 应用函数校准标签
        df["标签"] = df["文本内容"].apply(correct_labels)
        
        # 提取并重命名所需的列
        cleaned_df = df[["文本内容", "标签"]].rename(columns={"文本内容": "text", "标签": "label"})
        
        # 保存清洗后的数据到新的Excel文件
        cleaned_df.to_excel(os.path.join(CLEAN_DATA_FOLDER, filename), index=False)

print("所有数据清洗完成，结果已保存到 clean_data 文件夹中。")
